/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description: TbeUnsortedSegmentSum的头文件
 * Author: zhangjian
 * Create: 2023-5-17
 */

#ifndef TBE_UNSORTED_SEGMENT_SUM_H
#define TBE_UNSORTED_SEGMENT_SUM_H

#include <string>
#include <unordered_map>
#include <nlohmann/json.hpp>
#include <hccl/hccl_types.h>
#include <mutex>
#include <map>
#include <list>
#include <hccl/base.h>
#include "tbe_vector_reduce.h"
#ifndef HCCD
#include "runtime/rt.h"
#include "graph/operator.h"
#include "graph/compute_graph.h"
#include "graph/node.h"
#endif
#include "op_tiling.h"

namespace TbeReduce {

class TbeUnsortedSegmentSum : public TbeVectorReduce {
public:
    TbeUnsortedSegmentSum();
    ~TbeUnsortedSegmentSum() override;
    HcclResult Init();
    HcclResult Run(const void *input, const void *segmentids,
        u64 inputCount, u64 outputConut, u32 valueDim, const HcclDataType dataType, void *stream, const void *dst);
private:
    void InitOpInfoMap910A();
    void InitOpInfoMap310P3();
    void InitOpInfoMap910B();
    HcclResult GetOpInfo(nlohmann::json &opDescInfo, nlohmann::json &opTilingInfo);
    HcclResult GetTilingRunInfo(nlohmann::json &opDescInfo, nlohmann::json &opTilingInfo,
        u64 inputCount, u64 outputCount, u32 valueDim, HcclDataType dataType, OpRunInfo &runInfo);
    HcclResult GetTilingRunInfo(nlohmann::json &opDescInfo,
        u64 inputCount, u64 outputCount, u32 valueDim, HcclDataType dataType, OpRunInfo &runInfo);
    HcclResult ExecuteKernelLaunch(const void *input, const void *segmentids, const void *num_segments,
        const void *dst, void *stream, const void *stubFunc, OpRunInfo &runInfo, void *tilingDataDevPtr,
        const void *overFlowPtr);
    HcclResult GetopInfoIndex(std::string &opInfoIndex);
    HcclResult GetCoreNum();

    std::map<std::string, nlohmann::json> opInfoMap_;
    std::mutex tilingDataMutex_;
    std::mutex stubFuncMutex_;
    bool isInit_{};
    void *overFlowPtr_ = nullptr;
    void *numSegments_ = nullptr;
};
} // namespace TbeReduce
#endif
